from PyQt5.QtCore import QSharedMemory
import pandas as pd
import pickle
import sys


class PandasSharedMemory:
    def __init__(self, key):
        """
        初始化共享内存管理器

        参数:
            key: 共享内存的唯一标识键
        """
        self.key = key
        self.shared_mem = QSharedMemory()
        self.shared_mem.setKey(self.key)
        print(f'初始化共享内存，key为: {self.key}')

    def write_dataframe(self, df):
        """
        将 Pandas DataFrame 写入共享内存

        参数:
            df: 要写入的 Pandas DataFrame

        返回:
            bool: 是否写入成功
        """
        # 序列化 DataFrame
        try:
            data = pickle.dumps(df)
        except Exception as e:
            self._show_error(f"Data serialization failed: {str(e)}")
            return False

        # 清理现有共享内存
        self._cleanup()

        print(f'{self.key} isAttached:{self.shared_mem.isAttached()}')
        print(len(data))
        # 创建新的共享内存段
        if not self.shared_mem.create(len(data)):
            self._show_error(f"Failed to create shared memory: {self.shared_mem.errorString()}")
            return False

        # 写入数据
        self.shared_mem.lock()
        try:
            buffer = self.shared_mem.data()
            buffer[:len(data)] = data
            print(f'DataFrame {self.key} successfully written to shared memory')
            return True
        except Exception as e:
            self._show_error(f"Failed write to shared memory: {str(e)}")
            return False
        finally:
            self.shared_mem.unlock()

    def read_dataframe(self):
        """
        从共享内存读取 Pandas DataFrame

        返回:
            DataFrame: 读取到的 DataFrame 或 None (如果失败)
        """
        # 附加到共享内存
        if not self.shared_mem.attach():
            self._show_error(f"Failed read to attach to shared memory: {self.shared_mem.errorString()}")
            return None

        # 读取数据
        self.shared_mem.lock()
        try:
            buffer = self.shared_mem.data()
            # 获取所有数据 (QSharedMemory 没有直接的大小获取方法)
            # 这里假设整个内存都是我们的数据
            data = bytes(buffer)
            # 反序列化
            try:
                df = pickle.loads(data)
                if not isinstance(df, pd.DataFrame):
                    self._show_error("Data is not a Pandas DataFrame")
                    return None
                return df
            except Exception as e:
                self._show_error(f"Data deserialization failed: {str(e)}")
                return None
        finally:
            self.shared_mem.unlock()

    def modify_dataframe(self, new_df):
        """
        从共享内存读取 Pandas DataFrame

        返回:
            DataFrame: 读取到的 DataFrame 或 None (如果失败)
        """

        # 序列化 DataFrame
        try:
            data = pickle.dumps(new_df)
        except Exception as e:
            self._show_error(f"Data serialization failed: {str(e)}")
            return False
        # 附加到共享内存
        if not self.shared_mem.isAttached():
            self._show_error(f"Failed modify to attach to shared memory: {self.shared_mem.errorString()}")
            return None

        self.shared_mem.lock()
        try:
            buffer = self.shared_mem.data()
            buffer[:len(data)] = data
            return True
        except Exception as e:
            self._show_error(f"Failed to write to shared memory: {str(e)}")
            return False
        finally:
            self.shared_mem.unlock()

    def _cleanup(self):
        """清理共享内存资源"""
        if self.shared_mem.isAttached():
            self.shared_mem.detach()

    @staticmethod
    def _show_error(message):
        """显示错误消息"""
        # QMessageBox.critical(None, "Shared Memory Error", message)
        print(f"Shared Memory Error：{message}")

    # def __del__(self):
    #     """析构函数确保资源释放"""
    #     self._cleanup()
